#include <png++/png.hpp>
#include <GL/gl.h>
#include <GL/glu.h>
#include <sstream>
#include <iomanip>
#include <iostream>

using namespace std;

void write_img(unsigned char* pixels, int w, int h, const std::string& prefix, int n) {
    png::image< png::rgb_pixel > img(w, h);
    for (int j = 0; j < h; ++j) {
        for (int i = 0; i < w; ++i) {
            int coord = j*w*3+i*3;
            img[j][i] = png::rgb_pixel(pixels[coord+0],pixels[coord+1],pixels[coord+2]);
        }
    }
    stringstream formatter;
    formatter << prefix;
    formatter << setw(5) << setfill('0') << n << ".png";
    
    img.write(formatter.str());
}

void screenshot(int w, int h, const std::string& prefix, int n) {
    unsigned char* pixels = new unsigned char[w*h*3];
    if (pixels == NULL) throw std::exception();
    glPixelStorei(GL_PACK_ALIGNMENT, 1);
    glReadPixels(0, 0, w, h, GL_RGB, GL_UNSIGNED_BYTE, pixels);

    write_img(pixels, w, h, prefix, n);

    delete [] pixels;
}

